import pymysql
import datetime
import time

class mysqlIndustryClass():
    def __init__(self):
        # mysql数据库
        db_config = {
            'host': '8.141.23.4',  # 数据库服务器地址
            'port': 3306,  # 数据库端口
            'user': 'remote_user',  # 数据库用户名
            'password': 'mjquPH3kGKFYsGykH&x',  # 数据库密码
            'db': 'business_data'  # 数据库名
        }
        self.connection = pymysql.connect(**db_config)

    # 获取企业网页提取的信息
    def getWebInfoByCompanyName(self, companyName):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL查询语句
                sql_query = """
                  SELECT website, is_official_website, product, business, industry
                  FROM ent_industry_class
                  WHERE company_name = %s
                  """
                # 执行查询
                cursor.execute(sql_query, (companyName,))
                # 获取查询结果
                result = cursor.fetchone()
                if result:
                    res = [v if v is not None else "" for v in result]
                else:
                    res = [""]*5
                return res
        except Exception as e:
            print(f"[ERROR]mysqlIndustryClass.getWebInfoByCompanyName,{companyName}: {e}")
            return [""]*5

    # 更新企业网页的信息
    def setWebInfoByCompanyName(self, companyName, creditCode, website, isOfficialWebsite, product, business, industry):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL更新或插入语句
                sql_query = """
                  INSERT INTO ent_industry_class (company_name, credit_code, website, is_official_website, product, business, industry)
                  VALUES (%s, %s, %s, %s, %s, %s, %s)
                  ON DUPLICATE KEY UPDATE
                    credit_code=VALUES(credit_code),
                    website=VALUES(website),
                    is_official_website=VALUES(is_official_website),
                    product=VALUES(product),
                    business=VALUES(business),
                    industry=VALUES(industry),
                    update_time=CURRENT_TIMESTAMP
                  """
                # 执行查询
                cursor.execute(sql_query, (companyName, creditCode, website, isOfficialWebsite, product, business, industry))
                # 提交事务
                self.connection.commit()
        except Exception as e:
            # 如果发生错误则回滚
            self.connection.rollback()
            print(f"[ERROR]mysqlIndustryClass.setWebInfoByCompanyName,{companyName}: {e}")

    # 更新企业的打标信息
    def setIndustryClassByCompanyName(self, companyName, industryClass1, industryClass2, industryClass3, explain, certainty, version):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL更新或插入语句
                sql_query = """
                  INSERT INTO ent_industry_class (company_name, industry_class_first, industry_class_second, industry_class_third, `explain`, certainty, `version`, update_time)
                  VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                  ON DUPLICATE KEY UPDATE
                    industry_class_first=VALUES(industry_class_first),
                    industry_class_second=VALUES(industry_class_second),
                    industry_class_third=VALUES(industry_class_third),
                    `explain`=VALUES(`explain`),
                    certainty=VALUES(certainty),
                    `version`=VALUES(`version`),
                    update_time=VALUES(update_time);
                """
                # 执行查询，添加current_time作为最后一个参数
                cursor.execute(sql_query, (
                    companyName, industryClass1, industryClass2, industryClass3, explain, certainty, version,
                    datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
                # 提交事务
                self.connection.commit()
        except Exception as e:
            # 如果发生错误则回滚
            self.connection.rollback()
            print(f"[ERROR]mysqlIndustryClass.setIndustryClassByCompanyName,{companyName}: {e}")

    # 批量插入企业打标信息
    def batchSetIndustryClass(self, batch_data):
        """
        批量插入企业打标信息
        :param batch_data: 批量数据列表，每个元素为(company_name, class1, class2, class3, explain, certainty, version)
        :return: 是否成功
        """
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL语句
                sql_query = """
                  INSERT INTO ent_industry_class (
                    company_name,
                    industry_class_first,
                    industry_class_second,
                    industry_class_third,
                    `explain`,
                    certainty,
                    `version`,
                    update_time
                  ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                  ON DUPLICATE KEY UPDATE
                    industry_class_first=VALUES(industry_class_first),
                    industry_class_second=VALUES(industry_class_second),
                    industry_class_third=VALUES(industry_class_third),
                    `explain`=VALUES(`explain`),
                    certainty=VALUES(certainty),
                    `version`=VALUES(`version`),
                    update_time=VALUES(update_time);
                """

                # 准备数据，添加当前时间
                current_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                prepared_data = [(*data, current_time) for data in batch_data]

                # 开始事务
                self.connection.begin()

                # 执行批量插入
                cursor.executemany(sql_query, prepared_data)

                # 提交事务
                self.connection.commit()
                return True

        except Exception as e:
            # 回滚事务
            self.connection.rollback()
            print(f"[ERROR]mysqlIndustryClass.batchSetIndustryClass: {e}")
            return False

    # 获取企业打标信息
    def getIndustryClassByCompanyName(self, companyName):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL查询语句
                sql_query = """
                  SELECT industry_class_first, industry_class_second, industry_class_third, certainty, `explain`, `version`
                  FROM ent_industry_class
                  WHERE company_name = %s
                  """
                # 执行查询
                cursor.execute(sql_query, (companyName,))
                # 获取查询结果
                result = cursor.fetchone()
                if result:
                    res = [v if v is not None else "" for v in result]
                else:
                    res = [""] * 6
                return res
        except Exception as e:
            print(f"[ERROR]mysqlIndustryClass.getIndustryClassByCompanyName,{companyName}: {e}")
            return [""] * 6

    # 获取企业打标信息
    def getIndustryAnnotation(self, fullLabel):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL查询语句
                sql_query = """
                  SELECT annotation
                  FROM ent_industry_label
                  WHERE full_name = %s
                  """
                # 执行查询
                cursor.execute(sql_query, (fullLabel,))
                # 获取查询结果
                result = cursor.fetchone()
                if result:
                    res = result[0]
                else:
                    res = ""
                return res
        except Exception as e:
            print(f"[ERROR]mysqlIndustryClass.getIndustryAnnotation,{fullLabel}: {e}")
            return ""

    # 批量插入企业打标信息
    def batchSetIndustryClass(self, batch_data):
        """批量插入企业打标信息
        :param batch_data: 批量数据列表，每个元素为(company_name, class1, class2, class3, explain, certainty, version)元组
        :return: 是否成功
        """
        try:
            with self.connection.cursor() as cursor:
                # 构建批量插入SQL
                sql = """
                    INSERT INTO ent_industry_class
                    (company_name, industry_class_first, industry_class_second, industry_class_third,
                    `explain`, certainty, `version`, update_time)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                    ON DUPLICATE KEY UPDATE
                        industry_class_first=VALUES(industry_class_first),
                        industry_class_second=VALUES(industry_class_second),
                        industry_class_third=VALUES(industry_class_third),
                        `explain`=VALUES(`explain`),
                        certainty=VALUES(certainty),
                        `version`=VALUES(`version`),
                        update_time=VALUES(update_time);
                """

                # 准备批量数据
                now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                values = [(*data, now) for data in batch_data]

                # 执行批量插入
                cursor.executemany(sql, values)
                self.connection.commit()
                return True

        except Exception as e:
            # 如果发生错误则回滚
            self.connection.rollback()
            print(f"[ERROR]mysqlIndustryClass.batchSetIndustryClass: {e}")
            return False

    # 获取表结构
    def getTableCreateStatement(self, tableName):
        try:
            with self.connection.cursor() as cursor:
                # 构建SQL查询语句
                sql_query = f"SHOW CREATE TABLE `{tableName}`"
                # 执行查询
                cursor.execute(sql_query)
                # 获取查询结果
                result = cursor.fetchone()
                if result:
                    return result[1]  # 返回建表语句
                return ""
        except Exception as e:
            print(f"[ERROR]mysqlIndustryClass.getTableCreateStatement,{tableName}: {e}")
            return ""

if __name__ == "__main__":
    # 测试批量插入功能
    db = mysqlIndustryClass()

    # 准备测试数据
    test_data = [
        ("测试公司1", "一级分类1", "二级分类1", "三级分类1", "说明1", "高", "1.0"),
        ("测试公司2", "一级分类2", "二级分类2", "三级分类2", "说明2", "中", "1.0"),
        ("测试公司3", "一级分类3", "二级分类3", "三级分类3", "说明3", "低", "1.0")
    ]

    # 测试批量插入
    print("\n测试批量插入...")
    start_time = time.time()
    success = db.batchSetIndustryClass(test_data)
    end_time = time.time()
    print(f"批量插入结果: {'成功' if success else '失败'}")
    print(f"耗时: {end_time - start_time:.4f}秒")

    # 验证数据
    print("\n验证插入数据...")
    for company in ["测试公司1", "测试公司2", "测试公司3"]:
        result = db.getIndustryClassByCompanyName(company)
        print(f"{company} 数据: {result}")

    # 清理测试数据
    print("\n清理测试数据...")
    with db.connection.cursor() as cursor:
        for company in ["测试公司1", "测试公司2", "测试公司3"]:
            cursor.execute("DELETE FROM ent_industry_class WHERE company_name = %s", (company,))
        db.connection.commit()
